// See the file "COPYING" in the main distribution directory for copyright.

#include "zeek/packet_analysis/protocol/ieee802_11/IEEE802_11.h"

using namespace zeek::packet_analysis::IEEE802_11;

IEEE802_11Analyzer::IEEE802_11Analyzer() : zeek::packet_analysis::Analyzer("IEEE802_11") { }

bool IEEE802_11Analyzer::AnalyzePacket(size_t len, const uint8_t* data, Packet* packet)
	{
	u_char len_80211 = 24; // minimal length of data frames

	if ( len_80211 >= len )
		{
		Weird("truncated_802_11_header", packet);
		return false;
		}

	u_char fc_80211 = data[0]; // Frame Control field
	bool is_amsdu = false;

	// Skip non-data frame types (management & control).
	if ( ! ((fc_80211 >> 2) & 0x02) )
		return false;

	// Skip subtypes without data.
	if ( (fc_80211 >> 4) & 0x04 )
		return false;

	// 'To DS' and 'From DS' flags set indicate use of the 4th address field.
	if ( (data[1] & 0x03) == 0x03 )
		len_80211 += packet->L2_ADDR_LEN;

	if ( len_80211 >= len )
		{
		Weird("truncated_802_11_header", packet);
		return false;
		}

	// Look for the QoS indicator bit.
	if ( (fc_80211 >> 4) & 0x08 )
		{
		// Store off whether this is an A-MSDU header, which indicates that there are
		// multiple packets following the 802.11 header.
		is_amsdu = (data[len_80211] & 0x80) == 0x80;

		// Check for the protected bit. This means the data is encrypted and we can't
		// do anything with it.
		if ( data[1] & 0x40 )
			return true;

		len_80211 += 2;
		}

	if ( len_80211 >= len )
		{
		Weird("truncated_802_11_header", packet);
		return false;
		}

	// Determine link-layer addresses based on 'To DS' and 'From DS' flags
	switch ( data[1] & 0x03 )
		{
		case 0x00:
			packet->l2_src = data + 10;
			packet->l2_dst = data + 4;
			break;

		case 0x01:
			packet->l2_src = data + 10;
			packet->l2_dst = data + 16;
			break;

		case 0x02:
			packet->l2_src = data + 16;
			packet->l2_dst = data + 4;
			break;

		case 0x03:
			packet->l2_src = data + 24;
			packet->l2_dst = data + 16;
			break;
		}

	// skip the 802.11 data header
	data += len_80211;
	len -= len_80211;

	if ( ! is_amsdu )
		{
		return HandleInnerPacket(len, data, packet);
		}
	else
		{
		size_t amsdu_padding = 0;
		size_t encap_index = packet->encap ? packet->encap->Depth() : 0;

		while ( len > 0 )
			{
			if ( len < 14 )
				{
				Weird("truncated_802_11_amsdu_header", packet);
				return false;
				}

			// This is the length of everything after the A-MSDU subframe header.
			size_t amsdu_len = (data[12] << 8) + data[13];
			if ( len < amsdu_len )
				{
				Weird("truncated_802_11_amsdu_packet", packet);
				return false;
				}

			// Skip the A-MSDU subframe header. This should place us at the start of an LLC header.
			data += 14;
			len -= 14;

			if ( ! HandleInnerPacket(amsdu_len, data, packet) )
				{
				Weird("invalid_802_11_amsdu_inner_packet", packet);
				return false;
				}

			data += amsdu_len;
			len -= amsdu_len;

			// Each A-MSDU subframe is padded by up to 3 bytes to make a multiple of 4. This padding
			// isn't included in the length field value. The padding also doesn't happen with the
			// last subframe, so check to see that we can even subtract it. Unfortunately, there
			// isn't a frame counter in the header so we just have trust that it all works out.
			amsdu_padding = amsdu_len % 4;
			if ( len >= amsdu_padding )
				{
				data += amsdu_padding;
				len -= amsdu_padding;
				}

			// Pop encapsuations back up to the level where we started processing so that the next
			// subframe gets the same encapsulation stack.
			if ( packet->encap )
				{
				while ( packet->encap->Depth() > encap_index )
					packet->encap->Pop();
				}
			}

		return true;
		}
	}

bool IEEE802_11Analyzer::HandleInnerPacket(size_t len, const uint8_t* data, Packet* packet) const
	{
	// Make sure there's room for an LLC header.
	if ( len < 8 )
		{
		Weird("truncated_802_11_llc_header", packet);
		return false;
		}

	// Check that the DSAP and SSAP are both SNAP and that the control field indicates that this is
	// an unnumbered frame.  The organization code (24bits) needs to also be zero to indicate that
	// this is encapsulated ethernet.
	if ( data[0] == 0xAA && data[1] == 0xAA && data[2] == 0x03 && data[3] == 0 && data[4] == 0 &&
	     data[5] == 0 )
		{
		data += 6;
		len -= 6;
		}
	else
		{
		// If this is a logical link control frame without the possibility of having a protocol we
		// care about, we'll just skip it for now.
		return false;
		}

	// Get the protocol and skip the rest of the LLC header.
	uint32_t protocol = (data[0] << 8) + data[1];
	data += 2;
	len -= 2;

	return ForwardPacket(len, data, packet, protocol);
	}
